from .ai_query import AIQueryManager
from .ollama_query import OllamaQueryManager
from .db_manager import DatabaseManager

class NL2SQL:
    """
    NL2SQL 是一个自然语言转SQL并执行查询的统一接口类。
    支持通过不同AI模型（如DeepSeek、Ollama）将自然语言问题转为SQL，并查询数据库。
    适合集成到各类Python项目中，便于复用和扩展。
    """
    def __init__(self, model='ollama'):
        """
        初始化NL2SQL对象。
        参数：
            model (str): 指定使用的AI模型服务，'deepseek' 或 'ollama'，默认为'ollama'。
        """
        if model == 'deepseek':
            self.query_manager = AIQueryManager()
        else:
            self.query_manager = OllamaQueryManager()
        self.db = DatabaseManager()

    def nl2result(self, question: str) -> dict:
        """
        将自然语言问题转为SQL并执行查询。
        参数：
            question (str): 自然语言问题描述。
        返回：
            dict: 查询结果，包含SQL语句、执行结果或错误信息。
        """
        return self.query_manager.query(question)

    def nl2sql(self, question: str) -> str:
        """
        将自然语言问题转为SQL，仅返回处理后的SQL语句（不执行查询）。
        参数：
            question (str): 自然语言问题描述。
        返回：
            str: 处理后的、可直接执行的SQL语句。
        """
        # 获取AI返回后验证和处理后的SQL
        success, sql_or_err = self.query_manager.get_valid_sql(question)
        if success:
            print(f"处理后的SQL: {sql_or_err}")
            return sql_or_err
        else:
            return f"SQL验证失败: {sql_or_err}"


    def get_tables(self):
        """
        获取数据库中所有表的信息。
        返回：
            list: 表名及注释的列表。
        """
        return self.db.get_tables()

    def get_table_columns(self, table_name: str):
        """
        获取指定表的所有字段信息。
        参数：
            table_name (str): 表名。
        返回：
            list: 字段信息列表。
        """
        return self.db.get_table_columns(table_name)

    def get_table_data(self, table_name: str, limit: int = 10):
        """
        获取指定表的部分数据。
        参数：
            table_name (str): 表名。
            limit (int): 返回的最大行数，默认10。
        返回：
            list: 表数据（每行为字典）。
        """
        return self.db.get_table_data(table_name, limit)
